/*+**************************************************************************/
/***                                                                      ***/
/***   Copyright (C) by Dierk Ohlerich                                    ***/
/***   all rights reserved                                                ***/
/***                                                                      ***/
/***   To license this software, please contact the copyright holder.     ***/
/***                                                                      ***/
/**************************************************************************+*/

asc
{
  cbuffer Wz4SSAOShaderCamera : register(c0) : slot vs 0
  {
    row_major float4x4 MV;
    row_major float4x4 MVP;         // model view projection matrix

    extern void Set(const sViewport &view)
    {
      MV = view.ModelView;
      MVP = view.ModelScreen;
    }
  };
  
  cbuffer Wz4SSAOShaderPixel : register(c0) : slot ps 0
  {
    float4 posScale;
    float2 projZ;
    float3 projScale;
    float3 invProjScale;

    extern void Set(const sViewport &view)
    {
      posScale.x = 1.0f / view.Target.SizeX();
      posScale.y = 1.0f / view.Target.SizeY();
      posScale.z = (0.5f - view.Target.x0) * posScale.x;
      posScale.w = (0.5f - view.Target.y0) * posScale.y;
      
      projScale.x =  view.ZoomX / 2.0f;
      projScale.y = -view.ZoomY / 2.0f;
      projScale.z = 1.0f;
      
      invProjScale.x =  2.0f / view.ZoomX;
      invProjScale.y = -2.0f / view.ZoomY;
      invProjScale.z = 1.0f;
      
      projZ.x = view.Proj.k.z;
      projZ.y = view.Proj.l.z;
    }
  };
  
  cbuffer Wz4SSAOShaderRandom : register(c4) : slot ps 1
  {
    float4 sampleVec[32];
  };
}

/****************************************************************************/

material Wz4SSAODepthShader
{
  prepare
  {
    VertexShader = VS();
    PixelShader = PS();
  }
  
  vs
  {
    asc vs_3_0
    {
      use Wz4SSAOShaderCamera;
      void main
      (
        in float4 in_pos : POSITION,
        
        out float4 out_pos : POSITION,
        out float2 out_zw : TEXCOORD0,
      )
      {
        out_pos = mul(in_pos,MVP);
        out_zw = out_pos.zw;
      }
    }
  }
  
  ps
  {
    asc ps_3_0
    {
      void main
      (
        in float2 zw : TEXCOORD0,
        out float4 result : COLOR0
      )
      {
        result = zw.x / zw.y;
      }
    }
  }
};

material Wz4SSAONormalShader
{
  prepare
  {
    VertexShader = VS();
    PixelShader = PS();
  }
  
  vs
  {
    asc vs_3_0
    {
      use Wz4SSAOShaderCamera;
      void main
      (
        in float4 in_pos : POSITION,
        
        out float4 out_pos : POSITION,
      )
      {
        out_pos = in_pos;
      }
    }
  }
  
  ps
  {
    asc ps_3_0
    {
      use Wz4SSAOShaderPixel;
    
      sampler2D sDepth : register(s0);
      
      float3 reconstruct(float2 pos)
      {
        float depth = tex2D(sDepth,pos).x;
        float z = projZ.y / (depth - projZ.x);
        return invProjScale * float3(pos.xy-0.5,1.0) * z;
      }
    
      void main
      (
        in float2 posxy : VPOS,
        out float4 result : COLOR0,
      )
      {
        posxy = posxy * posScale.xy + posScale.zw;
        float3 mid = reconstruct(posxy);
        float3 rgt = reconstruct(posxy+float2(posScale.x,0));
        float3 upp = reconstruct(posxy+float2(0,posScale.y));
        
        float3 normal = normalize(cross(rgt-mid, upp-mid));
        result = float4((normal + 1) * 0.5,0.0);
      }
    }
  }
};

/****************************************************************************/

material Wz4SSAOShader
{
  prepare
  {
    VertexShader = VS();
    PixelShader = PS();
  }

  vs
  {
    asc vs_3_0
    {
      use Wz4SSAOShaderCamera;
      void main
      (
        in float4 in_pos : POSITION,
        in float3 in_norm  : NORMAL,
        
        out float4 out_pos : POSITION,
        out float4 out_norm : TEXCOORD0,
      )
      {
        out_norm.xyz  = mul(float4(in_norm,0),MV).xyz;
        out_pos = mul(in_pos,MVP);
        out_norm.w = out_pos.z;
      }
    }
  }

  ps
  {
    asc ps_3_0
    {
      use Wz4SSAOShaderPixel;
      use Wz4SSAOShaderRandom;
      
      sampler2D sDepth : register(s0);
      sampler2D sRandom : register(s1);
      
      static const float sampleRadius = 0.75;
      
      static const float minOcclude = 1.0/128.0;
      static const float maxOcclude = 24.0;
      
      static const int nSamples = 16;
      
      float unprojectZ(float depth)
      {
        return projZ.y / (depth - projZ.x);
      }
      
      float3 reconstruct(float2 pos)
      {
        return invProjScale * float3(pos.xy-0.5,1.0) * unprojectZ(tex2D(sDepth, pos).r);
      }
      
      float3 reconstruct2(float2 pos, float z)
      {
        return invProjScale * float3(pos.xy-0.5,1.0) * z;
      }
      
      float2 project(float3 pos)
      {
        return projScale.xy * (pos.xy / pos.z) + 0.5;
      }
      
      float occlusionFunction(float x)
      {
        /*x = saturate((x - minOcclude) / (maxOcclude - minOcclude));
        if(x > 0)
        {
          x = 1.0 - x;
          x *= x;
          x *= x;
        }*/
        if(x <= minOcclude)
          x = 0.0;
        else
          x = saturate(1.0 / (x*x) - 1.0 / (maxOcclude*maxOcclude));
        
        return x;
      }
    
      void main
      (
        in float4 normal : TEXCOORD0,
        in float2 posxy : VPOS,
        
        out float4 result : COLOR0,
      )
      {
        float2 pos01 = posxy * posScale.xy + posScale.zw;
        float3 pos = reconstruct2(pos01, normal.w);
        float3 random = tex2D(sRandom,posxy.xy / 4).xyz;
        float totalOcclude = 0;
              
        for(int i=0;i<nSamples;i++)
        {
          float3 sampleDir = sampleVec[i].xyz;
          sampleDir = reflect(random, sampleDir);
          if(dot(sampleDir,normal.xyz) < 0)
            sampleDir = -sampleDir;
          
          // pick occlusion sample
          float3 samplePoint = pos + sampleRadius * sampleDir;
          float depth = tex2D(sDepth, project(samplePoint)).r;
          float actualZ = unprojectZ(depth);
          totalOcclude += occlusionFunction(samplePoint.z - actualZ);
        }
        
        result = (1 - totalOcclude / nSamples) * saturate(-normalize(normal.xyz).z);
      }
    }
  }
};

/****************************************************************************/

material Wz4SSAOBlurShader
{
  prepare
  {
    VertexShader = VS();
    PixelShader = PS();
  }

  vs
  {
    asc vs_3_0
    {
      void main
      (
        in float4 in_pos : POSITION,
        out float4 out_pos : POSITION,
      )
      {
        out_pos = in_pos;
      }
    }
  }

  ps
  {
    asc ps_3_0
    {
      //use Wz4SSAOShaderPixel;
      
      uniform float4 posScale : register(c0);
      sampler2D sAO : register(s0);
      
      static const float3 samples[9] =
      {
        // corners
        { -1.50,-1.50, 4.0/25.0, },
        {  1.50,-1.50, 4.0/25.0, },
        { -1.50, 1.50, 4.0/25.0, },
        {  1.50, 1.50, 4.0/25.0, },
        // edges
        {  0.00,-1.50, 2.0/25.0, },
        { -1.50, 0.00, 2.0/25.0, },
        {  1.50, 0.00, 2.0/25.0, },
        {  0.00, 1.50, 2.0/25.0, },
         // center
        {  0.00, 0.00, 1.0/25.0, },
        
        /*// corners
        -1.25,-1.25,25.0/256.0,
         1.25,-1.25,25.0/256.0,
        -1.25, 1.25,25.0/256.0,
         1.25, 1.25,25.0/256.0,
        // edges
         0.00,-1.20,30.0/256.0,
        -1.20, 0.00,30.0/256.0,
         1.20, 0.00,30.0/256.0,
         0.00, 1.20,30.0/256.0,
        // center
         0.00, 0.00,36.0/256.0,*/
      };
      
      void main
      (
        in float4 normal : TEXCOORD0,
        in float2 posxy : VPOS,
        
        out float4 result : COLOR0
      )
      {
        result = 0;
        for(int i=0;i<9;i++)
          result += samples[i].z * tex2D(sAO,(posxy+samples[i].xy) * posScale.xy + posScale.zw);
      }
    }
  }
};
